
//% color="#FF6347" iconWidth=50 iconHeight=40
namespace cnnlib{

    export function initmodel() {}
    //% block="初始化模块" blockType="command"
    export function init(parameter: any, block: any) {

        Generator.addImport(`from CNNDetection import *
from DynamicVs import *
from Draw import *
from CNNtrain import *
`);
        Generator.addDeclaration("fire_de",`
model_st = None
fire_de = ImageProcessor()
ImagePro = ImagePro()
visua = DynamicVisualizer()
CNNtrain = CNNtrain()`)

    }
    //% block="加载图像路径[PATH]结果存入变量[VAR]" blockType="command"
    //% PATH.shadow="string" PATH.defl="image.jpg"
    //% VAR.shadow="variables_get" VAR.defl="img"
    export function loadimage(parameter: any, block: any) {
        let path=parameter.PATH.code;
        let var_name=parameter.VAR.code;        

        Generator.addCode(`${var_name}=ImagePro.load_image(${path},model_st)
img_tensor = ${var_name}`)

    }
    //% block="黑白边界检测" blockType="tag"
    export function initmode1l() {}
    //% block="黑白边界检测，白色区域的比例[NUM]%" blockType="command"

    //% NUM.shadow="range"   NUM.params.min=0    NUM.params.max=100   NUM.defl=60
    export function readcap(parameter: any, block: any) {
        let num=parameter.NUM.code;
        Generator.addCode(`fire_de.generate_bw_image(${num}/100)`)
    }
   //% block="卷积处理" blockType="tag"
   export function initmode2l() {}

    //% block="设置卷积核类型[KERNEL]结果存入变量[VAR]" blockType="command"
    //% KERNEL.shadow="dropdown" KERNEL.options="KERNEL"
    //% VAR.shadow="normal" VAR.defl="kernel"
    export function setkernel(parameter: any, block: any) {
        let kernel=parameter.KERNEL.code;
        let var_name=parameter.VAR.code;
        
        Generator.addCode(`${var_name}=visua.set_kernel('${kernel}')
kernel_type = '${kernel}'`)
 
    }
    //% block="动态展示卷积核[VAR]对图像[IMG]执行卷积操作的过程，间隔时间[TIME]" blockType="command"
    //% VAR.shadow="variables_get" VAR.defl="kernel"
    //% IMG.shadow="normal" IMG.defl="img"
    //% TIME.shadow="range"  TIME.params.min=10    TIME.params.max=1000   TIME.defl=10
    export function convolve(parameter: any, block: any) {
        let kernel=parameter.VAR.code;
        let img=parameter.IMG.code;
        let time=parameter.TIME.code;
        Generator.addDeclaration("fire_de1",`model_st = 'gray'`)
        Generator.addCode(`
visua.show_animation(${img},${time})`)
    }
   //% block="卷积、池化处理" blockType="tag"
    export function initmode334() {}

   export function initmode3l() {}
  
    //% block="利用卷积核[HH]对图像[IMG]执行卷积操作，结果存入变量[VAR]" blockType="command"
    //% IMG.shadow="normal" IMG.defl="img"
    //% HH.shadow="normal" HH.defl="kernel"
    //% VAR.shadow="variables_get" VAR.defl="feature"
    export function apply_convolution(parameter: any, block: any) {
        let img=parameter.IMG.code;
        let hh=parameter.HH.code;
        let var_name=parameter.VAR.code
        Generator.addDeclaration("fire_de1",`model_st = 'tensor'`)
        Generator.addInit("ImagePro",`pool_type = ''`)
        Generator.addCode(`
kernel = torch.from_numpy(${hh}).unsqueeze(0).unsqueeze(0).float()  
kernel = ImagePro.expand_kernel(kernel) 
pool_type = ''
${var_name} = ImagePro.apply_convolution(${img}, kernel)
`)
    }
    
    //% block="对卷积处理后得到的特征图[IMG]，执行[TYPE]池化操作，结果存入变量[VAR]" blockType="command"
    //% IMG.shadow="normal" IMG.defl="feature"
    //% TYPE.shadow="dropdown" TYPE.options="TYPE"
    //% VAR.shadow="normal" VAR.defl="feature_pooled"
    export function apply_pooling(parameter: any, block: any) {
        let img=parameter.IMG.code;
        let type=parameter.TYPE.code;
        let var_name=parameter.VAR.code
        Generator.addCode(`pool_type = '${type}'
${var_name} = ImagePro.apply_pooling(${img},'${type}')
`)
     
    }
    
    //% block="显示原始图像（标题[TIQ]），显示特征图[IMG]（标题[TITLE]），字号[NUM]" blockType="command"
    //% TIQ.shadow="string" TIQ.defl="原始图像"
    //% IMG.shadow="normal" IMG.defl="feature"
    //% TITLE.shadow="string" TITLE.defl="特征图"
    //% NUM.shadow="number" NUM.defl=6
    export function draw(parameter: any, block: any) {
        let tiq=parameter.TIQ.code;
        let img=parameter.IMG.code;
        let title=parameter.TITLE.code;
        let num=parameter.NUM.code;
        Generator.addCode(`ImagePro.draw(img_tensor,${img},kernel_type,${tiq},${title},${num},pool_type)`)
     
    }

    //% block="加载图像路径[PATH]，将图像的 RGB（红、绿、蓝）三个颜色通道分离并显示" blockType="command"
    //% PATH.shadow="string" PATH.defl="image.jpg"
    export function loadimage2(parameter: any, block: any) {
        let path=parameter.PATH.code;
        Generator.addCode(`ImagePro.draw_rgb(${path})`)
    }


    //% block="训练模型" blockType="tag"
    export function initmode4l() {}
    
    //% block="加载训练集路径[PATHA]测试集路径[PATHB]，得到训练集数据[TRAIN]与测试集数据[TEST]" blockType="command"
    //% PATHA.shadow="string" PATHA.defl="Train"
    //% PATHB.shadow="string" PATHB.defl="Test"
    //% TRAIN.shadow="normal" TRAIN.defl="train_data_loader"
    //% TEST.shadow="normal" TEST.defl="test_data_loader"
    export function loaddata(parameter: any, block: any) {
        let patha=parameter.PATHA.code;
        let pathb=parameter.PATHB.code;
        let train=parameter.TRAIN.code;
        let test=parameter.TEST.code;
        Generator.addCode(`train_data_loader, test_data_loader = CNNtrain.load_data(${patha},${pathb})`)
    }
    
    //% block="训练卷积神经网络模型，训练集数据[A]测试集数据[B]训练轮数[C]" blockType="command"
    //% A.shadow="normal" A.defl="train_data_loader"
    //% B.shadow="normal" B.defl="test_data_loader"
    //% C.shadow="number" C.defl=3
    export function trainmodel(parameter: any, block: any) {
        let a=parameter.A.code;
        let b=parameter.B.code;
        let c=parameter.C.code;
        Generator.addInit('model',`save_path= "model.pt"`,true)
        Generator.addCode(`model_st,model_value = CNNtrain.train(${a},${b},${c},save_path)`)
    }

    //% block="模型操作" blockType="tag"
    export function initmode5l1() {}
    //% block="计算模型准确度" blockType="reporter"

    export function testmodel(parameter: any, block: any) {

        Generator.addCode(`model_value`)
    }
     //% block="---" blockType="tag"
     export function initmode5l11() {}
    //% block="保存模型[PATH]" blockType="command"
    //% PATH.shadow="string" PATH.defl="model.pt"
    export function savemodel(parameter: any, block: any) {
        let path=parameter.PATH.code;
        Generator.addInit('model',`save_path= ${path}`,true)
    }

    //% block="加载模型[PATH]结果存入变量[MODEL]" blockType="command"
    //% PATH.shadow="string" PATH.defl="model.pt"
    //% MODEL.shadow="variables_get" MODEL.defl="model"
    export function loadmodel(parameter: any, block: any) {
        let path=parameter.PATH.code;
        let model=parameter.MODEL.code;
        Generator.addCode(`
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

checkpoint = torch.load(${path}, map_location=device)
if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
    model = models.resnet50(pretrained=False)
    model.fc = nn.Linear(model.fc.in_features, checkpoint['num_classes'])
    model.load_state_dict(checkpoint['model_state_dict'])
  
    CNNtrain.CLS2ID = checkpoint['CLS2ID']
    CNNtrain.ID2CLS = checkpoint['ID2CLS']
    CNNtrain.num_classes = checkpoint['num_classes']
else:
    model = checkpoint

${model} = model.to(device)`)
    }
    //% block="模型推理" blockType="tag"
    export function initmode5l114() {}


    //% block="显示图像[IMG] 标题[TITLE]字号[NUM]" blockType="command"
    //% IMG.shadow="normal" IMG.defl="img"
    //% TITLE.shadow="string" TITLE.defl="原始图像"
    //% NUM.shadow="number" NUM.defl=20
    export function detect(parameter: any, block: any) {
        let img=parameter.IMG.code;
        let title=parameter.TITLE.code;
        let num=parameter.NUM.code;
        Generator.addCode(`CNNtrain.image_display(${img},${title},${num})`)
    }
    //% block="调用卷积神经网络模型[MODEL]对图像[IMG]进行分类" blockType="command"
    //% MODEL.shadow="normal" MODEL.defl="model"
    //% IMG.shadow="normal" IMG.defl="img"
    export function detectfire(parameter: any, block: any) {
        let model=parameter.MODEL.code;
        let img=parameter.IMG.code;
        Generator.addCode(`confidence_value, class_value=CNNtrain.model_inference_results(${img},${model})`)
    }
    //% block="预测分类结果[ID]" blockType="reporter"
    //% ID.shadow="dropdown" ID.options="class_value"
    export function showresult(parameter: any, block: any) {
        let id=parameter.ID.code;
        Generator.addCode(`${id}`)
}